import numpy as np
import os
from tqdm import tqdm
import json
# from pypcd import pypcd
#import pcl
import argparse
import threading
import shutil


# kitti ppt
# class dz dy dx x y z r

def mkdir(path):
    if not os.path.exists(path):
        os.mkdir(path)


# def save_ori_bin(dataname):
#     bin_path = "./{}_ppt/training/innovusion/".format(dataname)


# srcfile 需要复制、移动的文件
# dstpath 目的地址
def mycopyfile(srcfile, dataname, bag_name):  # 复制函数
    dstpath = "./{}_ppt/training/innovusion/{}_".format(dataname, bag_name)
    if not os.path.isfile(srcfile):
        print("%s not exist!" % (srcfile))
    else:
        fpath, fname = os.path.split(srcfile)  # 分离文件名和路径
        # if not os.path.exists(dstpath):
        #   os.makedirs(dstpath)  # 创建路径
        shutil.copy(srcfile, dstpath + fname)  # 复制文件
        print("copy %s -> %s" % (srcfile, dstpath + fname))


# def pcd2bin(dataname, pcd_path, bag_name):
#     pcd_file_name = os.path.basename(pcd_path)
#     bin_path = "./{}_ppt/training/innovusion/{}_{}".format(dataname, bag_name, pcd_file_name.replace(".pcd", ".bin"))
#
#     points = pcl.load_XYZI(pcd_path).to_array()[:, :4]  # 读取pcd文件
#
#     # points = pypcd.PointCloud.from_path(pcd_path)
#     # np_x = (np.array(points.pc_data['x'], dtype=np.float32)).astype(np.float32)
#     # np_y = (np.array(points.pc_data['y'], dtype=np.float32)).astype(np.float32)
#     # np_z = (np.array(points.pc_data['z'], dtype=np.float32)).astype(np.float32)
#     # np_i = (np.array(points.pc_data['intensity'], dtype=np.float32)).astype(np.float32) / 256
#     np_x = (np.array(points[:, 0], dtype=np.float32)).astype(np.float32)
#     np_y = (np.array(points[:, 1], dtype=np.float32)).astype(np.float32)
#     np_z = (np.array(points[:, 2], dtype=np.float32)).astype(np.float32)
#     np_i = (np.array(points[:, 3], dtype=np.float32)).astype(np.float32)
#
#     points_bin = np.transpose(np.vstack((np_x, np_y, np_z, np_i)))
#     points_bin.tofile(bin_path)
#     return points


def loadlabel(label_path):
    if not os.path.exists(label_path):
        return
    with open(label_path) as f:
        labels = json.load(f)
        objs = []
        for label in labels:
            obj = {}
            position = label['psr']['position']
            rotation = label['psr']['rotation']['z']
            scale = label['psr']['scale']
            x, y, z = position["x"], position["y"], position["z"]
            xs, ys, zs = scale["x"], scale["y"], scale["z"]
            obj["obj_type"], obj["obj_id"] = label['obj_type'], label['obj_id']
            obj["box"] = [x, y, z, xs, ys, zs, rotation]
            objs.append(obj)
    return objs


class_map = {"Car": "car", "Pedestrian": "person", "Truck": "truck",
             "Bus": "bus", "Motorcycle": "motorbike", "MotorcyleRider": "rider",
            "Bicycle": "bicycle", "BicycleRider": "rider", "Van": "van",
             "Excavator": "engineering", "Crane": "engineering","TricycleRider":"rider"}


#class_map = {"Car": "car", "Pedestrian": "person", "Truck": "truck",
#             "Bus": "bus", "Motorcycle": "rider", "MotorcyleRider": "rider",
#             "Bicycle": "rider", "BicycleRider": "rider", "Van": "bus",
#             "Excavator": "truck", "Crane": "truck", "TricycleRider":"rider"}


#class_map = {"Truck": "Truck", "Excavator": "Excavator", "Car": "Car",
#             "Wall": "Wall", "Pedestrian": "Pedestrian", "Stone": "Stone",
#             "Dust": "Dust", "Mound": "Mound", "Obstacle": "Obstacle",
#             }

def save_label(dataname, objs, label_path, bag_name):
    file_name = os.path.basename(label_path).replace(".json", ".txt")
    label_path = "./{}_ppt/training/label_2/{}_{}".format(dataname, bag_name, file_name)
    with open(label_path, "w") as f:
        for obj in objs:
            x, y, z, dx, dy, dz, r = obj["box"]
            # 修改类
            if (obj["obj_type"] == 0 or obj["obj_type"] == "rider"):
                continue
            # elif (obj["obj_type"] == 'BicycleRider' or obj["obj_type"] == 'MotorcyleRider')::
            #     c = 'rider'
            # else:
            c = class_map[obj["obj_type"]]
            s = "{} {} {} {} {} {} {} {}".format(c, dz, dy, dx, x, y, z, r)
            if not obj is objs[-1]:
                s = s + "\n"
            f.write(s)

    pass


def init_dir(dataname):
    mkdir("./{}_ppt/".format(dataname))
    mkdir("./{}_ppt/training/".format(dataname))
    mkdir("./{}_ppt/training/innovusion/".format(dataname))
    mkdir("./{}_ppt/training/label_2/".format(dataname))
    # os.system("ln -s ./{}_ppt/training/ ./{}_ppt/testing/".format(dataname,dataname))


# def run(bar, bag_idx, bags_name, lidar_dir, dataname, arg):
def run(bar, bag_idx, bags_name, lidar_dir, dataname):
    for lidar_file in bar:
        bar.set_description("bag/all {} {}/{}".format(bag_name, bag_idx + 1, len(bags_name)))
        pcd_path = os.path.join(lidar_dir, lidar_file)
        # label_path = pcd_path.replace("/lidar/", "/label_filter/").replace(".pcd", ".json")
        label_path = pcd_path.replace("/lidar/", "/label/").replace(".bin", ".json")
        if not os.path.exists(os.path.dirname(label_path)):
            break
        # pcd2bin(dataname, pcd_path, bag_name)
        mycopyfile(pcd_path, dataname, bag_name)
        objs = loadlabel(label_path)
        if(objs is None):
            continue
        save_label(dataname, objs, label_path, bag_name)


args = argparse.ArgumentParser()
args.add_argument("--thread", action='store_true', default=False)
arg = args.parse_args()

dataname = "8014"
bags_name = os.listdir(dataname)
init_dir(dataname)
for bag_idx, bag_name in enumerate(bags_name):
    bag_dir = os.path.join(dataname, bag_name)
    lidar_dir = "{}/lidar/".format(bag_dir)
    if os.path.isdir(lidar_dir):
        lidar_files = os.listdir(lidar_dir)
        labels_files = os.listdir(lidar_dir.replace("/lidar/", "/label/"))
        bar = tqdm(lidar_files)
        run(bar, bag_idx, bags_name, lidar_dir, dataname)
    # if arg.thread:
    #     threading.Thread(target=run,
    #                      args=(bar, bag_idx, bags_name, lidar_dir, dataname, arg)).start()

